from docx2python import docx2python
from docx2python.iterators import iter_paragraphs, iter_tables
from docx import Document
from docx.shared import Pt, RGBColor, Inches
from docx.enum.text import WD_PARAGRAPH_ALIGNMENT, WD_ALIGN_PARAGRAPH
from docx.enum.table import WD_TABLE_ALIGNMENT
from docx.oxml.ns import qn
from docx.oxml import parse_xml
import os
import re
import math

class DocxProcessingError(Exception):
    """自定义文档处理异常基类"""
    pass

class DocxHandler:
    @staticmethod
    def read_docx_content(file_path, content_types=None, keep_formatting=True, max_length=None):
        """
        改进版文档解析方法，使用文档流解析方案
        """
        try:
            # 参数校验
            if not os.path.exists(file_path):
                raise FileNotFoundError(f"文件不存在: {file_path}")
                
            # 设置默认参数
            content_types = content_types or ['text', 'image', 'table']
            
            with docx2python(file_path) as doc_result:
                ordered_elements = []
                element_counter = 0
                
                # 按文档流顺序处理元素
                for element in doc_result.document:
                    element_counter += 1
                    
                    # 处理段落（包含文本和图片）
                    if element.tag.endswith('p'):
                        if 'text' in content_types or 'image' in content_types:
                            DocxHandler._process_paragraph(
                                element, doc_result, ordered_elements,
                                content_types, keep_formatting, max_length
                            )
                    
                    # 处理表格
                    elif element.tag.endswith('tbl'):
                        if 'table' in content_types:
                            DocxHandler._process_table(
                                element, ordered_elements,
                                keep_formatting, max_length
                            )
                
                return ordered_elements
                
        except Exception as e:
            raise DocxProcessingError(f"文档解析失败: {str(e)}") from e

    @staticmethod
    def _process_paragraph(element, doc_result, elements, content_types, keep_formatting, max_length):
        """处理段落元素"""
        # 提取段落级格式
        para_format = DocxHandler._extract_paragraph_format(element)
        
        # 检查是否包含图片
        drawing = element.find('.//wp:inline', namespaces=doc_result.nsmap)
        if drawing and 'image' in content_types:
            DocxHandler._extract_image(element, doc_result, elements, para_format)
            return
            
        # 处理文本内容
        if 'text' not in content_types:
            return
            
        text_spans = []
        current_text = []
        current_format = None
        
        # 遍历段落内所有文本运行(run)
        for run in element.xpath('.//w:r', namespaces=doc_result.nsmap):
            run_text = ''.join(run.xpath('.//w:t/text()', namespaces=doc_result.nsmap))
            if not run_text.strip():
                continue
                
            run_format = DocxHandler._extract_run_format(run) if keep_formatting else {}
            
            # 合并相同格式的连续文本
            if run_format == current_format:
                current_text.append(run_text)
            else:
                if current_text:
                    text_spans.append({
                        'text': ''.join(current_text),
                        'formatting': current_format
                    })
                current_text = [run_text]
                current_format = run_format
                
        if current_text:
            text_spans.append({
                'text': ''.join(current_text),
                'formatting': current_format
            })
            
        # 构建文本元素
        if text_spans:
            full_text = ''.join(span['text'] for span in text_spans)
            chunks = DocxHandler._split_content(
                content=text_spans,
                max_length=max_length,
                is_text=True
            )
            
            for chunk in chunks:
                elements.append({
                    'type': 'text',
                    'content': chunk if keep_formatting else full_text,
                    'formatting': para_format,
                    'position': len(elements) + 1
                })

    @staticmethod
    def _extract_paragraph_format(element):
        """提取段落级格式"""
        formatting = {}
        
        # 对齐方式
        align = element.xpath('./w:pPr/w:jc/@w:val', namespaces=element.nsmap)
        if align:
            formatting['alignment'] = align[0]
            
        # 缩进设置
        ind = element.xpath('./w:pPr/w:ind', namespaces=element.nsmap)
        if ind:
            ind = ind[0]
            formatting['indent'] = {
                'left': ind.get('w:left'),
                'right': ind.get('w:right'),
                'firstLine': ind.get('w:firstLine')
            }
            
        # 间距设置
        spacing = element.xpath('./w:pPr/w:spacing', namespaces=element.nsmap)
        if spacing:
            spacing = spacing[0]
            formatting['spacing'] = {
                'before': spacing.get('w:before'),
                'after': spacing.get('w:after'),
                'line': spacing.get('w:line')
            }
            
        return formatting

    @staticmethod
    def _extract_run_format(run):
        """提取文本运行格式"""
        format_dict = {}
        
        # 字体格式
        rPr = run.find('./w:rPr', namespaces=run.nsmap)
        if rPr is not None:
            # 字体名称
            font = rPr.xpath('./w:rFonts/@w:ascii', namespaces=run.nsmap)
            if font:
                format_dict['font'] = font[0]
                
            # 字号
            size = rPr.xpath('./w:sz/@w:val', namespaces=run.nsmap)
            if size:
                format_dict['size'] = f"{int(size[0])/2}pt"
                
            # 颜色
            color = rPr.xpath('./w:color/@w:val', namespaces=run.nsmap)
            if color and color[0] != 'auto':
                format_dict['color'] = f"#{color[0]}"
                
            # 样式
            format_dict['bold'] = len(rPr.xpath('./w:b')) > 0
            format_dict['italic'] = len(rPr.xpath('./w:i')) > 0
            format_dict['underline'] = len(rPr.xpath('./w:u')) > 0
            
        return format_dict

    @staticmethod
    def _extract_image(element, doc_result, elements, para_format):
        """提取图片及其格式"""
        blip = element.xpath('.//a:blip/@r:embed', namespaces=doc_result.nsmap)
        if not blip:
            return
            
        image_id = blip[0]
        image_data = doc_result.images.get(image_id)
        if not image_data:
            return
            
        # 提取图片格式
        img_format = {}
        inline = element.find('.//wp:inline', namespaces=doc_result.nsmap)
        if inline is not None:
            # 图片尺寸
            extents = inline.find('.//wp:extent', namespaces=doc_result.nsmap)
            if extents is not None:
                img_format['width'] = f"{Inches(int(extents.get('cx'))/914400)}in"
                img_format['height'] = f"{Inches(int(extents.get('cy'))/914400)}in"
                
            # 环绕方式
            wrap = inline.find('.//wp:wrap', namespaces=doc_result.nsmap)
            if wrap is not None:
                img_format['wrap'] = wrap.get('type')
                
        # 合并段落对齐方式
        img_format.update(para_format)
        
        elements.append({
            'type': 'image',
            'id': image_id,
            'data': image_data,
            'formatting': img_format,
            'position': len(elements) + 1
        })

    @staticmethod
    def _process_table(element, elements, keep_formatting, max_length):
        """处理表格元素"""
        try:
            table_data = []
            table_format = DocxHandler._extract_table_format(element)
            
            # 遍历表格行
            for row in element.xpath('./w:tr', namespaces=element.nsmap):
                row_data = []
                
                # 遍历单元格
                for cell in row.xpath('./w:tc', namespaces=element.nsmap):
                    cell_format = DocxHandler._extract_cell_format(cell)
                    cell_text = []
                    
                    # 提取单元格内容
                    for para in cell.xpath('.//w:p', namespaces=element.nsmap):
                        para_text = []
                        for run in para.xpath('.//w:r', namespaces=element.nsmap):
                            run_text = ''.join(run.xpath('.//w:t/text()', namespaces=element.nsmap))
                            if run_text.strip():
                                if keep_formatting:
                                    run_format = DocxHandler._extract_run_format(run)
                                    para_text.append({'text': run_text, 'formatting': run_format})
                                else:
                                    para_text.append(run_text)
                                    
                        if keep_formatting:
                            row_data.append({
                                'text': ''.join([t['text'] for t in para_text]),
                                'formatting': cell_format,
                                'spans': para_text
                            })
                        else:
                            row_data.append(''.join(para_text))
                            
                table_data.append(row_data)
            
            # 表格分段处理
            chunks = DocxHandler._split_content(
                content=table_data,
                max_length=max_length,
                is_text=False
            )
            
            for i, chunk in enumerate(chunks):
                elements.append({
                    'type': 'table',
                    'content': chunk,
                    'formatting': table_format if i == 0 else None,
                    'position': len(elements) + 1
                })
                
        except Exception as e:
            raise DocxProcessingError(f"表格处理失败: {str(e)}") from e

    @staticmethod
    def _extract_table_format(element):
        """提取表格级格式"""
        formatting = {}
        
        # 表格对齐
        jc = element.xpath('./w:tblPr/w:jc/@w:val', namespaces=element.nsmap)
        if jc:
            formatting['alignment'] = jc[0]
            
        # 表格边框
        tblBorders = element.xpath('./w:tblPr/w:tblBorders', namespaces=element.nsmap)
        if tblBorders:
            borders = {}
            for side in ['top', 'left', 'bottom', 'right']:
                border = tblBorders[0].xpath(f'./w:{side}', namespaces=element.nsmap)
                if border:
                    borders[side] = {
                        'color': border[0].get('w:color'),
                        'size': border[0].get('w:sz')
                    }
            formatting['borders'] = borders
            
        return formatting

    @staticmethod
    def _extract_cell_format(cell):
        """提取单元格格式"""
        formatting = {}
        
        # 单元格对齐
        jc = cell.xpath('./w:p/w:pPr/w:jc/@w:val', namespaces=cell.nsmap)
        if jc:
            formatting['alignment'] = jc[0]
            
        # 单元格背景色
        shading = cell.xpath('./w:tcPr/w:shd/@w:fill', namespaces=cell.nsmap)
        if shading:
            formatting['background'] = f"#{shading[0]}"
            
        return formatting

    @staticmethod
    def _split_content(content, max_length, is_text):
        """智能内容分段方法"""
        if not max_length:
            return [content]
            
        if is_text:
            # 文本分段：按字符长度分割
            full_text = ''.join([span['text'] for span in content])
            if len(full_text) <= max_length:
                return [content]
                
            chunks = []
            current_chunk = []
            current_length = 0
            
            for span in content:
                span_text = span['text']
                if current_length + len(span_text) > max_length:
                    # 需要拆分当前span
                    split_pos = max_length - current_length
                    chunks.append(current_chunk + [{
                        **span,
                        'text': span_text[:split_pos]
                    }])
                    current_chunk = [{
                        **span,
                        'text': span_text[split_pos:]
                    }]
                    current_length = len(current_chunk[0]['text'])
                else:
                    current_chunk.append(span)
                    current_length += len(span_text)
                    
            if current_chunk:
                chunks.append(current_chunk)
                
            return chunks
        else:
            # 表格分段：按行数分割
            total_rows = len(content)
            if total_rows == 0:
                return []
                
            avg_rows = max(1, math.floor(max_length / (len(content[0]) if content else 1)))
            return [content[i:i+avg_rows] for i in range(0, total_rows, avg_rows)]